from PyQt5 import QtCore, QtWidgets
import sys
from datetime import datetime
import cv2
# import mainwindow
# import Image
import paddlex as pdx
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtCore import *
from PyQt5.QtCore import Qt
from PyQt5.QtGui import *
from PyQt5.QtGui import QPixmap
from PyQt5.QtWidgets import *
from PyQt5.QtWidgets import QMessageBox, QFileDialog
from PyQt5.QtWidgets import QStyle

# 修改模型所在位置
predictor = pdx.deploy.Predictor('D:\\file4\\inference_model')
# path=None
# cap = cv2.VideoCapture(path)


class Ui_Form(object):
    def setupUi(self, Form):
        Form.setObjectName("Form")
        Form.resize(1588, 1043)
        self.label = QtWidgets.QLabel(Form)
        self.label.setGeometry(QtCore.QRect(30, 20, 1481, 751))
        self.label.setStyleSheet("background-color: rgb(0, 0, 0);")
        self.label.setText("")
        self.label.setObjectName("label")
        self.pushButton = QtWidgets.QPushButton(Form)
        self.pushButton.setGeometry(QtCore.QRect(300, 950, 131, 71))
        self.pushButton.setIconSize(QtCore.QSize(30, 30))
        self.pushButton.setObjectName("pushButton")
        self.textBrowser = QtWidgets.QTextBrowser(Form)
        self.textBrowser.setGeometry(QtCore.QRect(590, 790, 650, 231))
        self.textBrowser.setObjectName("textBrowser")
        self.pushButton_2 = QtWidgets.QPushButton(Form)
        self.pushButton_2.setGeometry(QtCore.QRect(300, 790, 131, 71))
        self.pushButton_2.setObjectName("pushButton_2")
        self.pushButton_3 = QtWidgets.QPushButton(Form)
        self.pushButton_3.setGeometry(QtCore.QRect(300, 870, 131, 71))
        self.pushButton_3.setObjectName("pushButton_3")
        self.pushButton.raise_()
        self.label.raise_()
        self.textBrowser.raise_()
        self.pushButton_2.raise_()
        self.pushButton_3.raise_()

        self.retranslateUi(Form)
        QtCore.QMetaObject.connectSlotsByName(Form)

    def retranslateUi(self, Form):
        _translate = QtCore.QCoreApplication.translate
        Form.setWindowTitle(_translate("Form", "火灾识别系统"))
        self.pushButton.setToolTip(_translate("Form", "<html><head/><body><p align=\"center\"><span style=\" font-size:12pt; font-weight:600;\">打开文件</span></p></body></html>"))
        self.pushButton.setWhatsThis(_translate("Form", "<html><head/><body><p align=\"center\"><span style=\" font-size:36pt; font-weight:600;\">打开文件</span></p></body></html>"))
        self.pushButton.setText(_translate("Form", "打开文件"))
        self.pushButton_2.setText(_translate("Form", "播放"))
        self.pushButton_3.setText(_translate("Form", "关闭"))




class Image_open(QMainWindow, Ui_Form):
    def __init__(self, parent=None):
        super(Image_open, self).__init__(parent)
        icon = QIcon('fire.ico')

        # Set the window icon
        self.setWindowIcon(icon)
        # UI界面
        self.setupUi(self)
        self.background()
        self.cap = cv2.VideoCapture()
        self.num = 1
        self.playing = False
        # 在label中播放视频
        self.init_timer()
    def background(self):
        # 文件选择按钮
        self.pushButton.clicked.connect(self.pre_judge)
        # 视频播放图标
        self.pushButton_2.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))  # 播放图标
        self.pushButton_3.setIcon(self.style().standardIcon(QStyle.SP_MediaStop))  # 停止图标

        # 用于开始播放视频的按钮
        self.pushButton_2.clicked.connect(self.play_file)  # 这是对应的函数
        # 用于关闭播放视频的按钮
        self.pushButton_3.clicked.connect(self.close_file)  # 这是对应的函数

        # 当播放的为图像时，设计这两个按钮不能点击，只有当播放的是视频时，才能点击
        self.pushButton_2.setEnabled(False)
        self.pushButton_3.setEnabled(False)

    def pre_judge(self):
            # 创建文件对话框，如果是视频，令self.video = True，如果是图像，令self.video = False，
            # 当self.video = None时，报错。
            self.video = None
            self.img_path = QFileDialog.getOpenFileName()[0]
            # self.textBrowser.setText(self.img_path)
            # self.textBrowser.setText("当前时间:"+"检测到小火\n")
            path=self.img_path
            video_type = [".mp4", ".mkv", ".MOV", "avi"]
            img_type = [".bmp", ".jpg", ".png", ".gif"]
            for vdi in video_type:
                if vdi not in self.img_path:
                    continue
                else:
                    self.video = True
                    # 当是视频时，将开始按钮置为可点击状态
                    self.pushButton_2.setEnabled(True)
            for ig in img_type:
                if ig not in self.img_path:
                    continue
                else:
                    self.video = False
                    img = QPixmap(self.img_path)
                    w = img.width()
                    h = img.height()
                    ratio = max(w / self.label.width(), h / self.label.height())
                    img.setDevicePixelRatio(ratio)
                    self.label.setAlignment(Qt.AlignCenter)
                    self.label.setPixmap(img)
            if self.video is None:
                QMessageBox.information(self, "警告", "我们暂时不支持此格式的文件！", QMessageBox.Ok)

    def open_file(self):
        self.cap = cv2.VideoCapture(self.img_path)
        if self.cap.isOpened():
            self.pushButton_2.setEnabled(True)
            self.pushButton_3.setEnabled(True)
            self.playing = True
            self.num = 0
            self.init_timer()
            self.timer.start(60)  # 设置视频帧率
            self.set_state()
        else:
            QMessageBox.information(self, "警告", "视频打开失败！", QMessageBox.Ok)



        # 打开本地视频文件
    def play_file(self):
            self.label.setEnabled(True)
            # 如果播放视频，则使得关闭视频按钮可用
            self.pushButton_3.setEnabled(True)
            # 视频流阻塞信号关闭
            self.timer.blockSignals(False)
            # 如果计时器没激活，证明是暂停阶段，需要重新播放，并把self.playing = True。
            if self.timer.isActive() is False:
                self.cap.open(self.img_path)
                self.timer.start(60)
                self.playing = True
                # 更换播放按钮为暂停按钮
                self.set_state()
            # 如果计时器激活了，并且num为奇数，证明是播放阶段，需要暂停播放，并把self.playing = False。
            elif self.timer.isActive() is True and self.num % 2 == 1:
                self.timer.blockSignals(True)
                self.playing = False
                self.num = self.num + 1
                self.set_state()
            # 如果计时器激活了，并且num为偶数，证明经过播放阶段，现在是暂停阶段，需要重新开始播放，并把self.playing = True。
            elif self.timer.isActive() is True and self.num % 2 == 0:
                self.num = self.num + 1
                self.timer.blockSignals(False)
                self.playing = True
                self.set_state()
            else:
                QMessageBox.information(self, "警告", "视频播放错误！", QMessageBox.Ok)

        # 关闭本地视频
    def close_file(self):
            self.cap.release()
            self.pushButton_2.setEnabled(False)
            self.pushButton_3.setEnabled(False)
            self.timer.stop()
            self.playing = False
            # 关闭视频将按钮置为可以播放
            self.img_path = ""
            self.textBrowser.setText(self.img_path)
            self.set_state()

        # 本地视频播放暂停转换图标按钮
    def set_state(self):
            if self.playing:
                # 暂停图标
                self.pushButton_2.setIcon(self.style().standardIcon(QStyle.SP_MediaPause))

            else:
                self.pushButton_2.setIcon(self.style().standardIcon(QStyle.SP_MediaPlay))

        # 播放视频画面
    def init_timer(self):
            self.timer = QTimer(self)
            self.timer.timeout.connect(self.show_pic)

        # 显示视频图像
    def show_pic(self):
            i = 0
            while self.cap.isOpened():
                now = datetime.now()
                ret, img = self.cap.read()
                if ret:
                    result = predictor.predict(img)
                    score = result[0]['score']
                    print(score)
                    if(i==0):
                        if score >= 0.3:
                            self.textBrowser.append("当前时间:"+now.strftime('%Y-%m-%d %H:%M:%S')+"系统检测到：火\n")
                    i=i+1
                    if (i == 9):
                        i = 0
                    vis_img = pdx.det.visualize(img, result, threshold=0.3, save_dir=None)
                    cur_frame = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
                    height, width = cur_frame.shape[:2]
                    pixmap = QImage(cur_frame, width, height, QImage.Format_RGB888)
                    pixmap = QPixmap.fromImage(pixmap)
                    ratio = max(width / self.label.width(), height / self.label.height())
                    pixmap.setDevicePixelRatio(ratio)
                    self.label.setAlignment(Qt.AlignCenter)
                    self.label.setPixmap(pixmap)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                else:
                    break

if __name__ == '__main__':
    app = QApplication(sys.argv)
    main = Image_open()
    main.show()
    sys.exit(app.exec_())